package main

import (
	"embed"
	"encoding/json"
	"fmt"
	"net/http"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"unsafe"

	"golang.org/x/sys/windows"
)

// 嵌入静态资源
//
//go:embed index.html
var fs embed.FS

// 用于存储每次查询的结果
type SearchResult struct {
	Addresses []uintptr `json:"addresses"`
	Total     int       `json:"total"`
}

// 存储每个 PID 的查询结果
var searchResults = make(map[uint32][]uintptr)
var resultsMutex sync.Mutex

// 获取所有正在运行的进程列表
func getProcessList() (map[string]int, error) {
	cmd := exec.Command("tasklist", "/NH", "/FO", "CSV")
	output, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	lines := strings.Split(string(output), "\r\n")
	processes := make(map[string]int)

	for _, line := range lines {
		if strings.TrimSpace(line) == "" {
			continue
		}
		fields := strings.Split(strings.Trim(line, "\""), "\",\"")
		if len(fields) >= 2 {
			exeName := fields[0]
			pid, err := strconv.Atoi(fields[1])
			if err == nil {
				processes[exeName] = pid
			}
		}
	}

	return processes, nil
}

// searchMemory 函数用于在指定进程的内存中搜索特定值。
// handle 是目标进程的句柄，通过该句柄可以对进程内存进行操作。
// value 是要搜索的整数值。
// prevAddresses 是上一次搜索得到的内存地址列表，如果为空，则进行全内存搜索；否则，仅在这些地址中搜索。
// 函数返回匹配该值的内存地址列表和可能出现的错误。
func searchMemory(handle windows.Handle, value int32, prevAddresses []uintptr) ([]uintptr, error) {
	// 用于存储本次搜索到的匹配值的内存地址
	var newAddresses []uintptr

	// 如果上一次搜索结果为空，进行全内存搜索
	if len(prevAddresses) == 0 {
		// 用于存储内存区域的基本信息
		var memInfo windows.MemoryBasicInformation
		// 从进程内存的起始地址开始搜索
		baseAddr := uintptr(0)

		// 遍历进程的整个内存空间
		for {
			// 获取当前内存区域的基本信息
			err := windows.VirtualQueryEx(handle, baseAddr, &memInfo, unsafe.Sizeof(memInfo))
			// 如果获取信息失败，说明已经遍历完所有内存区域，退出循环
			if err != nil {
				break
			}

			// 检查当前内存区域是否已提交且不是不可访问的页面
			if memInfo.State == windows.MEM_COMMIT && memInfo.Protect&windows.PAGE_NOACCESS == 0 {
				// 创建一个缓冲区，用于存储从内存中读取的数据
				buffer := make([]byte, memInfo.RegionSize)
				var nread uintptr
				// 从当前内存区域读取数据到缓冲区
				err := windows.ReadProcessMemory(handle, baseAddr, &buffer[0], memInfo.RegionSize, &nread)
				// 如果读取成功
				if err == nil {
					// 遍历缓冲区中的数据，每次读取 4 个字节（int32 类型）
					for i := 0; i < int(nread)-4; i++ {
						// 将缓冲区中的数据转换为 int32 类型
						val := *(*int32)(unsafe.Pointer(&buffer[i]))
						// 如果当前值等于要搜索的值
						if val == value {
							// 将该值所在的内存地址添加到结果列表中
							newAddresses = append(newAddresses, baseAddr+uintptr(i))
						}
					}
				}
			}

			// 移动到下一个内存区域
			baseAddr += memInfo.RegionSize
		}
	} else {
		// 如果上一次搜索结果不为空，仅在之前的地址列表中搜索
		for _, addr := range prevAddresses {
			// 创建一个 4 字节的缓冲区，用于存储从指定地址读取的数据
			buffer := make([]byte, 4)
			var nread uintptr
			// 从指定地址读取 4 个字节的数据到缓冲区
			err := windows.ReadProcessMemory(handle, addr, &buffer[0], 4, &nread)
			// 如果读取成功
			if err == nil {
				// 将缓冲区中的数据转换为 int32 类型
				val := *(*int32)(unsafe.Pointer(&buffer[0]))
				// 如果当前值等于要搜索的值
				if val == value {
					// 将该地址添加到结果列表中
					newAddresses = append(newAddresses, addr)
				}
			}
		}
	}

	// 返回本次搜索到的匹配值的内存地址列表，错误始终返回 nil
	return newAddresses, nil
}

// writeMemory 函数用于向指定进程的内存地址写入新的整数值。
// handle 是目标进程的句柄，通过该句柄可以对进程的内存进行操作。
// addr 是要写入的内存地址，类型为 uintptr。
// newValue 是要写入的新整数值，类型为 int32。
// 函数返回一个 error 类型的值，如果写入过程中出现错误，会返回具体的错误信息；若写入成功，则返回 nil。
func writeMemory(handle windows.Handle, addr uintptr, newValue int32) error {
	// nwrite 用于存储实际写入的字节数
	var nwrite uintptr
	// 调用 windows 系统的 WriteProcessMemory 函数，将 newValue 的值写入指定进程的内存地址 addr
	// (*byte)(unsafe.Pointer(&newValue)) 将 newValue 的指针转换为 *byte 类型，以便符合 WriteProcessMemory 函数的参数要求
	// 4 表示要写入的字节数，因为 int32 类型占 4 个字节
	err := windows.WriteProcessMemory(handle, addr, (*byte)(unsafe.Pointer(&newValue)), 4, &nwrite)
	return err
}

// 获取进程列表的 API 接口
func getProcessesHandler(w http.ResponseWriter, r *http.Request) {
	processes, err := getProcessList()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(processes)
}

// searchMemoryHandler 处理搜索内存的 HTTP 请求，根据传入的 PID、搜索值和页码返回相应的内存地址信息。
func searchMemoryHandler(w http.ResponseWriter, r *http.Request) {
	// 从请求的 URL 查询参数中获取 PID、搜索值和页码
	pidStr := r.URL.Query().Get("pid")
	valueStr := r.URL.Query().Get("value")
	pageStr := r.URL.Query().Get("page")

	// 将获取到的 PID 字符串转换为整数
	pid, err := strconv.Atoi(pidStr)
	if err != nil {
		http.Error(w, "Invalid PID", http.StatusBadRequest)
		return
	}

	// 将获取到的搜索值字符串转换为整数
	value, err := strconv.Atoi(valueStr)
	if err != nil {
		http.Error(w, "Invalid value", http.StatusBadRequest)
		return
	}

	// 默认页码为 1
	page := 1
	if pageStr != "" {
		page, err = strconv.Atoi(pageStr)
		if err != nil {
			http.Error(w, "Invalid page number", http.StatusBadRequest)
			return
		}
	}

	// 打开指定 PID 的进程，获取进程句柄，权限为读取进程内存
	handle, err := windows.OpenProcess(0x0010|0x0400, false, uint32(pid))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	// 函数结束时关闭进程句柄
	defer windows.CloseHandle(handle)

	// 加锁以安全访问全局的搜索结果
	resultsMutex.Lock()
	// 获取上一次该 PID 的搜索结果
	prevAddresses := searchResults[uint32(pid)]
	// 解锁
	resultsMutex.Unlock()

	// 调用 searchMemory 函数在进程内存中搜索指定值
	addresses, err := searchMemory(handle, int32(value), prevAddresses)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// 加锁以安全更新全局的搜索结果
	resultsMutex.Lock()
	// 更新该 PID 的搜索结果
	searchResults[uint32(pid)] = addresses
	// 解锁
	resultsMutex.Unlock()

	// 计算搜索结果的总数
	total := len(addresses)
	// 计算当前页的起始索引
	start := (page - 1) * 10
	// 计算当前页的结束索引
	end := start + 10
	if end > total {
		// 若结束索引超过总数，将其设置为总数
		end = total
	}

	// 创建响应对象，包含当前页的地址列表和总数
	response := SearchResult{
		Addresses: addresses[start:end],
		Total:     total,
	}

	// fmt.Println(response)

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(response)
}

// writeMemoryHandler 处理修改内存值的 HTTP 请求，根据传入的 PID、内存地址和新值，将新值写入指定进程的内存地址。
func writeMemoryHandler(w http.ResponseWriter, r *http.Request) {
	pidStr := r.URL.Query().Get("pid")
	addrStr := r.URL.Query().Get("addr")
	newValueStr := r.URL.Query().Get("newValue")

	pid, err := strconv.Atoi(pidStr)
	if err != nil {
		http.Error(w, "Invalid PID", http.StatusBadRequest)
		return
	}

	addr, err := strconv.ParseUint(addrStr, 0, 64)
	if err != nil {
		http.Error(w, "Invalid address", http.StatusBadRequest)
		return
	}

	newValue, err := strconv.Atoi(newValueStr)
	if err != nil {
		http.Error(w, "Invalid new value", http.StatusBadRequest)
		return
	}

	// 打开指定 PID 的进程，获取进程句柄，权限包括读取进程内存、写入进程内存
	handle, err := windows.OpenProcess(0x0010|0x0400|0x0020, false, uint32(pid))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer windows.CloseHandle(handle)

	// 调用 writeMemory 函数将新值写入指定的内存地址
	err = writeMemory(handle, uintptr(addr), int32(newValue))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
}

// 提供 HTML 页面
func indexHandler(w http.ResponseWriter, r *http.Request) {
	// http.ServeFile(w, r, "index.html")
	content, err := fs.ReadFile("index.html")
	if err != nil {
		http.Error(w, "Failed to read index.html", http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.Write(content)
}

// 清除查询结果的 API 接口
func resetSearchResultsHandler(w http.ResponseWriter, r *http.Request) {
	resultsMutex.Lock()
	defer resultsMutex.Unlock()
	searchResults = make(map[uint32][]uintptr)
	w.WriteHeader(http.StatusOK)
}

func main() {
	http.HandleFunc("/", indexHandler)
	http.HandleFunc("/processes", getProcessesHandler)
	http.HandleFunc("/search", searchMemoryHandler)
	http.HandleFunc("/write", writeMemoryHandler)
	http.HandleFunc("/reset", resetSearchResultsHandler)

	fmt.Println("Server started at :8081")
	http.ListenAndServe(":8081", nil)
}
